package schema

// Core types based on A2A spec (v0.3.0)
// Focused, minimal set to enable server/client scaffolding.

import (
    "encoding/json"
    "time"
)

// Role of a message author.
type Role string

const (
	RoleUser  Role = "user"
	RoleAgent Role = "agent"
)

// Part is the smallest unit of content in a Message or Artifact.
type Part interface{ isPart() }

// TextPart carries plain text.
type TextPart struct {
	Type string `json:"type"` // "text"
	Text string `json:"text"`
}

func (TextPart) isPart() {}

// FilePart references a file by URI (may include mimeType if known).
type FilePart struct {
	Type     string  `json:"type"` // "file"
	URI      string  `json:"uri"`
	MimeType *string `json:"mimeType,omitempty"`
}

func (FilePart) isPart() {}

// DataPart carries structured data as JSON.
type DataPart struct {
	Type string                 `json:"type"` // "data"
	Data map[string]interface{} `json:"data"`
}

func (DataPart) isPart() {}

// MarshalJSON for a generic Part slice uses type switch on concrete type.
func MarshalParts(parts []Part) ([]json.RawMessage, error) {
	out := make([]json.RawMessage, 0, len(parts))
	for _, p := range parts {
		var b []byte
		var err error
		switch v := p.(type) {
		case TextPart:
			b, err = json.Marshal(v)
		case FilePart:
			b, err = json.Marshal(v)
		case DataPart:
			b, err = json.Marshal(v)
		default:
			// unknown parts are encoded as null
			b = []byte("null")
		}
		if err != nil {
			return nil, err
		}
		out = append(out, b)
	}
	return out, nil
}

// Message is a communication turn with role and parts.
type Message struct {
	Role  Role   `json:"role"`
	Parts []Part `json:"-"`
	// PartsRaw is used for (un)marshalling.
	PartsRaw []json.RawMessage `json:"parts"`
}

func (m *Message) MarshalJSON() ([]byte, error) {
	type alias Message
	var tmp = struct {
		Role  Role              `json:"role"`
		Parts []json.RawMessage `json:"parts"`
	}{Role: m.Role}
	var err error
	tmp.Parts, err = MarshalParts(m.Parts)
	if err != nil {
		return nil, err
	}
	return json.Marshal(tmp)
}

// Artifact is an output generated by a task; composed of parts.
type Artifact struct {
	ID        string            `json:"id"`
	CreatedAt time.Time         `json:"createdAt"`
	Parts     []Part            `json:"-"`
	PartsRaw  []json.RawMessage `json:"parts"`
}

// TaskState enumerates lifecycle states per spec.
type TaskState string

const (
	TaskPending      TaskState = "pending"
	TaskRunning      TaskState = "running"
	TaskCompleted    TaskState = "completed"
	TaskFailed       TaskState = "failed"
	TaskCanceled     TaskState = "canceled"
	TaskBlocked      TaskState = "blocked"
	TaskAuthRequired TaskState = "auth-required"
)

// TaskStatus holds progress detail, optional message, and error fields.
type TaskStatus struct {
	State TaskState `json:"state"`
	// Optional human-readable message or structured data.
	Message   *DataPart `json:"message,omitempty"`
	Error     *string   `json:"error,omitempty"`
	UpdatedAt time.Time `json:"updatedAt"`
}

// Task is a unit of work in A2A.
type Task struct {
	ID        string     `json:"id"`
	ContextID *string    `json:"contextId,omitempty"`
	Status    TaskStatus `json:"status"`
	// Optional final output.
	Artifacts []Artifact `json:"artifacts,omitempty"`
}

// Push notification configuration for asynchronous updates.
type PushNotificationConfig struct {
	ID  string `json:"id"`
	URL string `json:"url"`
	// Opaque metadata for the server when sending webhooks.
	Secret *string `json:"secret,omitempty"`
}

// AgentCard describes an agent’s identity, capabilities, and endpoints.
type AgentCard struct {
    Name           string                 `json:"name"`
    Title          *string                `json:"title,omitempty"`
    Version        *string                `json:"version,omitempty"`
    Description    *string                `json:"description,omitempty"`
    Endpoints      map[string]string      `json:"endpoints,omitempty"` // e.g. {"jsonrpc":"/v1/jsonrpc"}
    Authentication map[string]interface{} `json:"authentication,omitempty"`
    Capabilities   []string               `json:"capabilities,omitempty"`
    capObj         *AgentCapabilities     `json:"-"`
}

// MarshalJSON supports dual-shape capabilities. If an object-shaped
// capabilities is available via the internal field, it is preferred.
func (a AgentCard) MarshalJSON() ([]byte, error) {
    type alias AgentCard
    // Build a transient struct with a flexible capabilities field
    var out = struct {
        Name           string                 `json:"name"`
        Title          *string                `json:"title,omitempty"`
        Version        *string                `json:"version,omitempty"`
        Description    *string                `json:"description,omitempty"`
        Endpoints      map[string]string      `json:"endpoints,omitempty"`
        Authentication map[string]interface{} `json:"authentication,omitempty"`
        Capabilities   interface{}            `json:"capabilities,omitempty"`
    }{
        Name:           a.Name,
        Title:          a.Title,
        Version:        a.Version,
        Description:    a.Description,
        Endpoints:      a.Endpoints,
        Authentication: a.Authentication,
    }
    // Prefer object-shaped capabilities if present
    if a.capObj != nil {
        out.Capabilities = a.capObj
    } else if len(a.Capabilities) > 0 {
        out.Capabilities = a.Capabilities
    }
    return json.Marshal(out)
}

// UnmarshalJSON accepts either the legacy []string or the spec object for capabilities.
func (a *AgentCard) UnmarshalJSON(b []byte) error {
    // Parse into a helper with raw capabilities
    var aux struct {
        Name           string                 `json:"name"`
        Title          *string                `json:"title,omitempty"`
        Version        *string                `json:"version,omitempty"`
        Description    *string                `json:"description,omitempty"`
        Endpoints      map[string]string      `json:"endpoints,omitempty"`
        Authentication map[string]interface{} `json:"authentication,omitempty"`
        Capabilities   json.RawMessage        `json:"capabilities,omitempty"`
    }
    if err := json.Unmarshal(b, &aux); err != nil {
        return err
    }
    a.Name = aux.Name
    a.Title = aux.Title
    a.Version = aux.Version
    a.Description = aux.Description
    a.Endpoints = aux.Endpoints
    a.Authentication = aux.Authentication
    // Default empty
    a.Capabilities = nil
    if len(aux.Capabilities) == 0 || string(aux.Capabilities) == "null" {
        return nil
    }
    // Try object first
    if len(aux.Capabilities) > 0 && aux.Capabilities[0] == '{' {
        var obj AgentCapabilities
        if err := json.Unmarshal(aux.Capabilities, &obj); err == nil {
            // Store object internally to keep API surface unchanged in phase 2
            a.capObj = &obj
            // Derive legacy list for compatibility
            var list []string
            if obj.Streaming != nil && *obj.Streaming {
                list = append(list, "streaming")
            }
            if obj.PushNotifications != nil && *obj.PushNotifications {
                list = append(list, "pushNotifications")
            }
            if obj.StateTransitionHistory != nil && *obj.StateTransitionHistory {
                list = append(list, "stateTransitionHistory")
            }
            a.Capabilities = list
            return nil
        }
        // fallthrough to try slice on error
    }
    // Try legacy []string
    var list []string
    if err := json.Unmarshal(aux.Capabilities, &list); err == nil {
        a.Capabilities = list
        return nil
    }
    // Unknown shape; ignore capabilities
    return nil
}
